# -*- coding: utf-8 -*-
"""
简化导入脚本 - 使用已验证的连接参数
"""

import subprocess
import os
import sys
import json
from pathlib import Path

project_root = Path(__file__).resolve().parent
sys.path.insert(0, str(project_root))

PSQL_PATH = r"D:\PostgreSQL\bin\psql.exe"
PASSWORD = "123456"

def run_sql(sql_command, database='postgres'):
    """执行 SQL 命令"""
    cmd = [
        PSQL_PATH,
        '-U', 'postgres',
        '-h', 'localhost',
        '-p', '5432',
        '-d', database,
        '-c', sql_command
    ]
    
    env = os.environ.copy()
    env['PGPASSWORD'] = PASSWORD
    
    result = subprocess.run(cmd, capture_output=True, env=env, timeout=30)
    return result.returncode == 0, result.stdout, result.stderr

def main():
    print("=" * 60)
    print("简化导入脚本")
    print("=" * 60)
    
    # 1. 创建数据库
    print("\n步骤1: 创建数据库...")
    run_sql("DROP DATABASE IF EXISTS QA_knowledge;")
    success, _, _ = run_sql("CREATE DATABASE QA_knowledge ENCODING 'UTF8';")
    
    if not success:
        print("✗ 创建数据库失败")
        return
    
    print("✓ 数据库创建成功")
    
    # 2. 创建表
    print("\n步骤2: 创建表...")
    
    create_table_sql = """
    CREATE TABLE qa_knowledge (
        question TEXT PRIMARY KEY,
        question_tokens JSONB,
        answer JSONB,
        context TEXT,
        category_id VARCHAR(50),
        created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
    );
    
    CREATE INDEX idx_qa_question_tokens_gin ON qa_knowledge USING GIN (question_tokens);
    CREATE INDEX idx_qa_tokens_array ON qa_knowledge USING GIN ((question_tokens->'tokens'));
    CREATE INDEX idx_qa_category_id ON qa_knowledge (category_id);
    """
    
    success, _, _ = run_sql(create_table_sql, database='QA_knowledge')
    
    if not success:
        print("✗ 创建表失败")
        return
    
    print("✓ 表创建成功")
    
    # 3. 准备数据
    print("\n步骤3: 准备数据...")
    
    from D_db.core.tokenizer import Tokenizer
    from D_db.io.excel_reader import ExcelReader
    
    excel_file = project_root / "C_filter" / "QA_final" / "大伙房水库调度规程.xlsx"
    
    excel_reader = ExcelReader()
    df = excel_reader.read_file(str(excel_file))
    qa_pairs = excel_reader.parse_qa_data(df)
    
    print(f"✓ 读取 {len(qa_pairs)} 个问答对")
    
    tokenizer = Tokenizer()
    
    # 4. 批量插入
    print("\n步骤4: 插入数据...")
    
    sql_file = project_root / "batch_insert.sql"
    
    with open(sql_file, 'w', encoding='utf-8') as f:
        for i, (question, answer, context) in enumerate(qa_pairs, 1):
            if i % 500 == 0:
                print(f"  准备 {i}/{len(qa_pairs)} 条记录...")
            
            token_data = tokenizer.tokenize_question(question)
            
            # 转义
            q = question.replace("'", "''").replace("\\", "\\\\")
            a = answer.replace("'", "''").replace("\\", "\\\\")
            c = (context.replace("'", "''").replace("\\", "\\\\")) if context else ''
            
            token_json = json.dumps(token_data, ensure_ascii=False).replace("'", "''")
            answer_json = json.dumps({"answer": answer}, ensure_ascii=False).replace("'", "''")
            
            sql = f"INSERT INTO qa_knowledge (question, question_tokens, answer, context, category_id) VALUES ('{q}', '{token_json}'::jsonb, '{answer_json}'::jsonb, '{c}', '21100150') ON CONFLICT (question) DO NOTHING;\n"
            f.write(sql)
    
    print(f"\n✓ SQL 文件已生成")
    
    # 5. 执行导入
    print("\n步骤5: 执行导入（这可能需要几分钟）...")
    
    cmd = [
        PSQL_PATH,
        '-U', 'postgres',
        '-h', 'localhost',
        '-p', '5432',
        '-d', 'postgre',
        '-f', str(sql_file)
    ]
    
    env = os.environ.copy()
    env['PGPASSWORD'] = PASSWORD
    
    result = subprocess.run(cmd, capture_output=True, env=env, timeout=300)
    
    if result.returncode == 0:
        print("✓ 导入成功")
        
        # 验证
        success, output, _ = run_sql("SELECT COUNT(*) FROM qa_knowledge;", database='QA_knowledge')
        if success:
            print(f"\n✓ 数据库中的记录数: {output.decode('utf-8', errors='ignore').strip()}")
        
        print("\n" + "=" * 60)
        print("✓ 完成！")
        print("=" * 60)
    else:
        print(f"✗ 导入失败")
        print(result.stderr.decode('utf-8', errors='ignore'))

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"错误: {e}")
        import traceback
        traceback.print_exc()
